# coding: utf-8
import pickle
from random import shuffle

import matplotlib.pyplot as plt
from tqdm import trange

from Encodings import CitySequence


def copy(target):
    return pickle.loads(pickle.dumps(target))


class GA:  # Genetic Algorithm
    def __init__(self, data):
        self.name = __class__.__name__
        # -- read instance from instance data --
        self.data = copy(data)
        self.V = self.data.cities  # city list
        self.n = self.data.n  # number of cities
        self.Depot = self.data.depot  # depot
        self.Distance = self.data.distance_matrix  # distance matrix
        # -- parameters ---
        self.NP = 10  # the number of population, must be 10x
        self.Pc = 0.8  # deepcopy probability
        # -- solution initialization --
        self.Population = [CitySequence().init_solution(cities=self.V, distance=self.Distance, style='greedy')
                           for _ in range(self.NP)]
        self.optima = copy(min(self.Population, key=lambda x: x.cost))
        self.evo_list = []
        # -- auxiliary --
        self.entropy = []  # the information entropy to measure the diversity of population

    def optimize(self, iteration=1000, trial=1):
        pbar = trange(iteration)
        pbar.set_description(f'{self.name} on {self.data.name} trial {trial}')
        for it in pbar:
            # update Pc value
            self.Pc = (self.Pc_max - self.Pc_min) * it / iteration + self.Pc_min
            self.Pc_list.append(self.Pc)  # record Pc value
            # perform the genetic operation of PGA
            shuffle(self.Population)
            for i in range(0, self.NP, 10):
                best = min(self.Population[i: i+10], key=lambda x: x.total)
                for idx in range(i, i + 10):
                    self.Population[idx] = best
                    if idx % 10 == 0:
                        pass
                    else:
                        # generate neighbor solutions according to different copy_prob
                        copy_prob = (self.Pc_max - self.Pc_min) * (idx % 10) / iteration + self.Pc_min
                        self.Population[idx] = self.Population[idx].neighborhood(copy_prob=copy_prob,
                                                                                 distance=self.Distance,
                                                                                 cost_func=self.tour_cost,
                                                                                 access=self.access)
            # update optima & record evolutionary history
            if self.optima.total > min(self.Population, key=lambda x: x.total).total:
                self.optima = copy(min(self.Population, key=lambda x: x.total))
            self.evo_list.append(self.optima.total)
            pbar.set_postfix_str(f'best: {self.optima.total:.2f}')
        return self


def edge_assembly_crossover(p1, p2):
    """ Edge Assembly Crossover (EAX) """
    
    return None


from TSP import TSP

data = TSP(path=f'../../data/TSPLIB/att48.tsp.txt')
optimizer = GA(data)

print(optimizer)
